import os
from pathlib import Path
import dashscope
from typing import Dict
import asyncio


class ImageAnalyzer:
    def __init__(
        self,
        api_key: str = os.getenv("DASHSCOPE_API_KEY") or "",
        model: str = "qwen-vl-max"
    ):
        """
        ImageAnalyzer 支持一次分析多张图片（整个 PDF 目录）
        """
        self.api_key = api_key
        self.model = model

    def analyze_pdf_directory(
        self,
        local_base_dir_raw: str | Path,
        base_url: str,
        prompt_text: str = "请根据这些图片描述整个 PDF 内容"
    ) -> Dict[str, str]:
        """
        批量分析整个 PDF 图片目录
        返回字典：{pdf_name: overall_description}
        """
        local_base_dir = Path(local_base_dir_raw)
        results = {}

        for pdf_dir in sorted(local_base_dir.iterdir()):
            if pdf_dir.is_dir():
                pdf_name = pdf_dir.name
                messages_content = []

                # 收集该 PDF 目录下所有图片的 URL
                for img_file in sorted(pdf_dir.iterdir()):
                    if img_file.is_file() and img_file.suffix.lower() in [".png", ".jpg", ".jpeg", ".webp"]:
                        url = f"{base_url}/{pdf_name}/{img_file.name}"
                        messages_content.append({"image": url})

                # 最后加上 prompt_text
                messages_content.append({"text": prompt_text})

                # 调用 DashScope 多模态模型
                response = dashscope.MultiModalConversation.call(
                    api_key=self.api_key,
                    model=self.model,
                    messages=[{"role": "user", "content": messages_content}]
                )

                # 直接取模型返回文本
                results[pdf_name] = response.get("content") or str(response)

        return results


if __name__ == '__main__':
    analyzer = ImageAnalyzer(api_key="sk-xxx")

    local_dir = Path(__file__).parent / "analysis"/"pdfs"
    if (not local_dir.exists()):
        print(f"{local_dir} not exists")
        quit(0)
    base_url = "https://static.21729901.xyz/res/data/analysis/pdfs"

    pdf_analysis_results = analyzer.analyze_pdf_directory(local_dir, base_url)

    for pdf_name, description in pdf_analysis_results.items():
        print(f"{pdf_name} 的分析结果：\n{description}\n")
